{-# LANGUAGE LambdaCase #-}
module Q.Typecheck
  where

--import Prelude hiding (lookup, pred)
import Control.Monad (unless)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State (StateT, state, get, runStateT)
import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Data.List (intersect, (\\), intercalate)
import Data.Map.Strict (Map)
import Data.Maybe (fromMaybe, maybeToList)
import Text.Printf (printf)

import qualified Data.Map.Strict as Map

import Q.IR.Common (Id(..))
import Q.IR.Surface

-- TODO:
-- checkSat needs to use some subset relation to check if the predicate is satisfied

--------------------------------------------------------------------------------
-- Symbolic values used by the typechecker

data Value
  = VTemp Temp
  | VInt Int
  | VBool Bool
  deriving (Eq, Ord)

newtype Temp = Temp Int
  deriving (Eq, Ord)

instance Show Value where
  show (VTemp t) = "_"
  show (VInt n)  = show n
  show (VBool True) = "true"
  show (VBool False) = "false"

instance Show Temp where
  show (Temp n) = printf "t%d" n

--------------------------------------------------------------------------------
-- Predicates

data Pred x
  = PBool
  | PInt
  | PIfBoolEq x Bool (Pred x)
  | PFun FunSig
  deriving (Eq)

type Constraints x = Map x [Pred x]

showPreds :: Show x => [Pred x] -> String
showPreds [] = "T"
showPreds ps = intercalate " /\\ " (show <$> ps)

showConstraints :: Show x => Constraints x -> String
showConstraints = intercalate ", " . map f . Map.toList
  where
    f (x, ps) = printf "%s: %s" (show x) (showPreds ps)

data FunSig = FunSig
  { fnArity :: Int
  , fnPre   :: Constraints Param
  , fnPost  :: Constraints Param }
  deriving (Eq)

data Param = RV | Arg Int
  deriving (Eq, Ord)

instance Show Param where
  show RV      = "rv"
  show (Arg i) = printf "a%d" i

instance Show x => Show (Pred x) where
  show PBool = "bool"
  show PInt  = "int"
  show (PIfBoolEq x b p) =
    printf "(%s | %s%s)"
      (show p)
      (if b then "" else "!")
      (show x)
  show (PFun funSig) =
    printf "fn(%s | %s) -> %s | %s"
      (intercalate ", " (show . Arg <$> [0..fnArity funSig - 1]))
      (showConstraints (fnPre funSig))
      (show RV)
      (showConstraints (fnPost funSig))

instPred :: (p -> Value) -> Pred p -> Maybe (Pred Temp)
instPred subst PBool      = Just PBool
instPred subst PInt       = Just PInt
instPred subst (PFun sig) = Just (PFun sig)
instPred subst (PIfBoolEq param b pred) =
  case subst param of
    VTemp x  -> PIfBoolEq x b <$> instPred subst pred
    VBool b' -> if b == b' then
                  instPred subst pred
                else
                  Nothing
    VInt _ ->
      Nothing

inst :: (p -> Value) -> Constraints p -> [(Value, Pred Temp)]
inst subst constraints = [ (subst param, pred')
                         | (param, preds) <- Map.toList constraints
                         , pred <- preds
                         , pred' <- maybeToList (instPred subst pred) ]

--------------------------------------------------------------------------------
-- Typechecker state monad

data S = S
  { counter :: Int
  , constraints :: Constraints Temp
  }

type TC a = StateT S (Except String) a

err :: String -> TC a
err msg = lift $ throwE msg

getConstraints :: TC (Constraints Temp)
getConstraints = constraints <$> get

gensym :: TC Int
gensym = state $ \s ->
  let s' = s { counter = counter s + 1 } in
  (counter s, s')

newTemp :: TC Temp
newTemp = Temp <$> gensym

predsOf :: Value -> TC [Pred Temp]
predsOf (VInt _)  = pure [PInt]
predsOf (VBool _) = pure [PBool]
predsOf (VTemp t) = fromMaybe [] . Map.lookup t <$> getConstraints

addConstraintsFor :: Temp -> [Pred Temp] -> TC ()
addConstraintsFor t ps = state $ \s ->
  let s' = s { constraints = Map.insertWith (++) t ps (constraints s) } in
  ((), s')

checkSat :: Value -> Pred Temp -> TC Bool
checkSat v p = elem p <$> predsOf v

--------------------------------------------------------------------------------
-- Typechecker

type Env = Map Id Value

tc :: Env -> Expr -> TC Value
tc env (EInt n)  = pure (VInt n)
tc env (EBool b) = pure (VBool b)

tc env (EVar x) = case Map.lookup x env of
  Just v  -> return v
  Nothing -> err $ printf "unbound: %s" (show x)

tc env (ELet x e1 e2) = do
  v <- tc env e1
  tc (Map.insert x v env) e2

tc env (EIf e0 e1 e2) = do
  v0  <- tc env e0
  sat <- checkSat v0 PBool
  unless sat (err "condition isn't a boolean")

  ps1 <- predsOf =<< tc env e1
  ps2 <- predsOf =<< tc env e2

  let
    meet = ps1 `intersect` ps2
    comb = meet ++
      (case v0 of
         VTemp cond  -> map (PIfBoolEq cond True)  (ps1 \\ meet) ++
                        map (PIfBoolEq cond False) (ps2 \\ meet)
         VBool True  -> ps1 \\ meet
         VBool False -> ps2 \\ meet
         VInt _      -> error "unreachable!")

  retT <- newTemp
  addConstraintsFor retT comb
  pure (VTemp retT)

tc env (ECall fId args) = do
  fp <- predsOf =<< tc env (EVar fId)
  fSig <- case [s | PFun s <- fp] of
            [s] -> pure s
            _   -> err $ printf "not a function: %s" (show fId)

  unless (fnArity fSig == length args)
    (err "wrong number of arguments")

  argVs <- mapM (tc env) args
  retT  <- newTemp

  let
    pre  = inst subst (fnPre fSig)
    post = inst subst (fnPost fSig)

    subst RV      = VTemp retT
    subst (Arg i) = argVs !! i

    chkCon (v,       p) = checkSat v p
    addCon (VTemp x, p) = addConstraintsFor x [p]
    addCon (_,       _) = pure ()

  sat <- and <$> mapM chkCon pre
  unless sat (err "function preconditions not satisifed")
  mapM_ addCon post
  pure (VTemp retT)

--------------------------------------------------------------------------------
-- Utilities for actually running the typechecker

op2 :: Pred Param -> Pred Param -> Pred Param -> FunSig
op2 p q r = FunSig
  { fnArity = 2
  , fnPre   = Map.fromList [(Arg 0, [p]), (Arg 1, [q])]
  , fnPost  = Map.fromList [(RV, [r])] }

pIf :: x -> Pred x -> Pred x -> [Pred x]
pIf x pTrue pFalse = [PIfBoolEq x True pTrue, PIfBoolEq x False pFalse]

prelude :: [(Id, FunSig)]
prelude =
  [ (Id "+", op2 PInt PInt PInt)
  , (Id "<", op2 PInt PInt PBool)
  , (Id "edgy", FunSig { fnArity = 2
                       , fnPre = Map.fromList
                         [ (Arg 0, [PBool])
                         , (Arg 1, pIf (Arg 0) PInt PBool) ]
                       , fnPost = Map.fromList
                         [ (RV, pIf (Arg 0) PBool PInt) ] })
  ]

initS :: [(Id, FunSig)] -> (S, Env)
initS funs = (s, env)
  where
    env = Map.fromList
          [ (funId, VTemp (Temp i))
          | ((funId, _), i) <- zip funs [0..] ]

    s = S { counter = length funs
          , constraints = Map.fromList
                          [ (Temp i, [PFun funSig])
                          | ((_, funSig), i) <- zip funs [0..] ] }

newtype TCResult = TCR (Either String [Pred Temp])

instance Show TCResult where
  show (TCR (Left err)) = "ERROR: " ++ err
  show (TCR (Right ps)) = showPreds ps

runTC :: S -> TC Value -> TCResult
runTC s0 t = TCR (fst <$> runExcept (runStateT t' s0))
  where
    t' = do
      v <- t
      predsOf v
      --pure (v, ps)

typecheck :: Expr -> TCResult
typecheck e = runTC s0 (tc env e)
  where
    (s0, env) = initS prelude


--------------------------------------------------------------------------------
-- Examples

-- >>> typecheck $ EBool True
-- bool
-- >>> typecheck $ EInt 0
-- int
-- >>> typecheck $ ELet (Id "z") (EInt 3) (EVar (Id "z"))
-- int
-- >>> typecheck $ EVar (Id "<")
-- fn(a0, a1 | a0: int, a1: int) -> rv | rv: bool
-- >>> typecheck $ EVar (Id "+")
-- fn(a0, a1 | a0: int, a1: int) -> rv | rv: int
-- >>> typecheck $ ECall (Id "+") [EInt 3, EInt 4]
-- int
-- >>> typecheck $ ECall (Id "<") [EInt 3, EInt 5]
-- bool
-- >>> typecheck $ ECall (Id "+") [EInt 3, EBool True]
-- ERROR: function preconditions not satisifed
-- >>> typecheck $ ECall (Id "edgy") [EBool True, EInt 4]
-- bool
-- >>> typecheck $ ECall (Id "edgy") [EBool False, EBool True]
-- int
-- >>> typecheck $ EIf (EBool True) (EInt 2) (EInt 4)
-- int
-- >>> typecheck $ EIf (EBool True) (EInt 2) (EBool False)
-- int
-- >>> typecheck $ EIf (EBool False) (EInt 2) (EBool False)
-- bool
-- >>> typecheck $ EIf (ECall (Id "<") [EInt 3, EInt 0]) (EInt 2) (EBool False)
-- (int | t3) /\ (bool | !t3)
-- >>> typecheck $ EIf (ECall (Id "<") [EInt 3, EInt 0]) (EInt 2) (EInt 0)
-- int
-- >>> typecheck $ EIf (EInt 0) (EInt 2) (EBool False)
-- ERROR: condition isn't a boolean
